import getopt
import sys
import oss2
import pymysql
import redis
import time
import json
import datetime
from PIL import Image
import io,os
import requests
import traceback
import base64
import random
from urllib.parse import urlparse

env = "prod" 

opts, args = getopt.getopt(sys.argv[1:], "","env=")
for opt, arg in opts:
    if opt == '--env':
        env = arg
        break

if env == "prod" :
    mysql_host = "rm-j6ctyd19275tjf7x1.mysql.rds.aliyuncs.com"
    mysql_user = "root"
    mysql_pwd = "qwemnb0011@"
    mysql_db = "aipainter"

    redis_host = "r-j6cmt5ayr2nq1m6zpp.redis.rds.aliyuncs.com"
    redis_db = 0
    redis_pwd = "r-j6cmt5ayr2nq1m6zpp:qpalmz1209@L"

    oss_endpoint = "oss-cn-hongkong-internal.aliyuncs.com"
else:
    mysql_host = "192.168.1.56"
    mysql_user = "root"
    mysql_pwd = "ubmy8920fe101f0aabffeeaecf1d9e04"
    mysql_db = "dl_oversea"

    redis_host = "192.168.1.56"
    redis_db = 0
    redis_pwd = "ubd83922c944c7e45b45ab5c86f98e7f"

    oss_endpoint = "oss-cn-hongkong.aliyuncs.com"

#mysql
conn = pymysql.connect(host=mysql_host,user=mysql_user, password=mysql_pwd, database=mysql_db)
cur = conn.cursor()

#redis
r = redis.StrictRedis(host=redis_host,port=6379, db=redis_db, password=redis_pwd)

#oss
ossAuth = oss2.Auth('LTAI5tMzkcDHweKw7762CzQn', '6cfufuTuoQpdwR7mckhxITwkLBWtcH')
ossBucket = oss2.Bucket(ossAuth, oss_endpoint, 'aipainter-overseas') 

def writeLog(data,machine):
    now = datetime.datetime.now().strftime('%Y-%m-%d')
    path = "../logs/{}_{}.txt".format(machine,now)
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    file = open(path,'a+')
    file.write(now+" "+data + '\n')
    file.close()

def getFileBase64(imgUrl):
    imgUrl += "!webp"
    response = requests.get(imgUrl)
    base64_data = base64.b64encode(response.content).decode('utf-8')
    return base64_data

def doTask(URL,machine) :
    queueKey = "painter:{}:queue".format(machine)
    
    loopNum = 0

    writeLog("loop start!!!",machine)

    while loopNum < 500:
        loopNum += 1

        sdinput = r.lpop(queueKey)
        if sdinput == None:
            writeLog("no task!!",machine)
            cur.execute('select 1')
            time.sleep(1)
            continue
        toStr = str(sdinput,encoding="utf-8")
        writeLog(toStr,machine)
        arr = json.loads(toStr)

        task_id = arr['task_id']
        writeLog("task_id : {} start!".format(task_id),machine)
    
        taskKey = "painter:task:{}".format(task_id)
        curUserId = arr['user_id'] if arr['user_id'] > 0 else arr["device_id"]
        userLenKey = "painter:user:len:{}".format(curUserId)
        pngFile = "{}.png".format(task_id)
        webpFile = "{}.webp".format(task_id)
        headers = {"User-Agent":arr["model"]}
        #drawing_image = arr['drawing_url']
        isHigh = arr['isHigh']
        style_id = arr["style_id"]

        #默认值drawing_image
        if "drawing_url" not in arr:
            drawing_image = ""
        else: 
            drawing_image = arr['drawing_url']

        #是否是延时任务
        delay_task = False
        if "is_dealy" in arr and arr["is_delay"] == 1:
            delay_task = True
        

        # 设置处理中状态
        r.hset(taskKey,'status',1)
        try :
            if arr['type'] == 1 or arr['type'] == 2:
                writeLog("任务类型：txt2img",machine)
                sampler_index = "Euler a" if arr["type"] == 1 else "DPM++ 2M Karras"
                payload = {
                    "prompt": arr["translator_prompt"],
                    "steps": 20,
                    "cfg_scale": "7.0",
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt": arr["negative_prompt"],
                    "sampler_index": sampler_index,
                }
                if isHigh == 1 :
                    hig_payload = {
                        "enable_hr":True,
                        "denoising_strength":0.5,
                        "hr_scale":2,
                        "hr_upscaler":"Latent"
                    }
                    payload.update(hig_payload)
                
                payload_json = json.dumps(payload)
            
                response = requests.post(url=f'{URL}/sdapi/v1/txt2img', data=payload_json, headers=headers, timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)

            else:
                writeLog("任务类型：img2img",machine)

                base64Data = getFileBase64(drawing_image)    

                rand_int = random.randint(0,1)
                interrogateModel = ""
                if rand_int == 0 : 
                    interrogateModel = "clip"
                else:
                    interrogateModel = "deepdanbooru"

                writeLog("interrogateModel：{}".format(interrogateModel),machine)

                payload = {
                    "image": base64Data,
                    "model": interrogateModel
                    #deepdanbooru,clip
                }

                payload_json_pro = json.dumps(payload)
                #writeLog("payload_json_pro：{}，header:{}".format(payload_json_pro,headers),machine)
                responsePro = requests.post(url=f'{URL}/sdapi/v1/interrogate', data=payload_json_pro, headers=headers, timeout=50).json()
                
                if "caption" not in responsePro:
                    responsePro['caption'] = ""
                arr["prompt"] = responsePro['caption']
                writeLog("图片解析的描述词：{}".format(responsePro['caption']),machine)
                
                base64dataArray = []
                base64dataArray.append(base64Data)

                payload = {
                    "init_images": base64dataArray,
                    "prompt" : responsePro['caption'],
                    "steps" : 20,
                    "denoising_strength": 0.65,
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt" : arr["negative_prompt"],
                    "sampler_index": "DPM++ 2M Karras",
                    "restore_faces": True
                }

                payload_json = json.dumps(payload)
                response = requests.post(url=f'{URL}/sdapi/v1/img2img', data=payload_json, headers=headers, timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)


            # png转webp Start
            #sCmd = "ffmpeg -nostdin -i " + pngFile + "  " + webpFile
            sCmd = "ffmpeg -nostdin -i " + pngFile + " -c:v libwebp -quality 90 -compression_level 6 " + webpFile
            writeLog("图片转换: {}".format(sCmd),machine)
            os.system(sCmd)
            writeLog("png转webp成功",machine)
            # png转webp End
            
            # 上传图片&删除本地图片
            path = "sd/opus/%s/%s" % (time.strftime("%Y-%m-%d", time.localtime()), webpFile)
            writeLog("本地图片地址：{}".format(webpFile),machine)
            writeLog("oss图片地址：{}".format(path),machine)
            ossBucket.put_object_from_file(path, webpFile)

            # 写入数据&更新状态
            url = 'https://overseas-files.dalleai.art/' + path

            sql = """
            insert into ap_opus(user_id, prompt, task_id, url, height, width, title, negative_prompt, type, drawing_id, drawing_image,is_new,device_id,is_high,dot,style_id,is_style_orginal,platform,machine_id) value(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            #writeLog(sql % (arr['user_id'], arr['prompt'], arr['task_id'], url, arr['height'], arr['width'], arr['prompt'], arr['negative_prompt'], arr['type'],arr['model_id'],drawing_image,1,arr["device_id"]),machine)
            cur.execute(sql, (arr['user_id'], arr['prompt'], arr['task_id'], url, arr['height'], arr['width'], arr['prompt'], arr['negative_prompt'], arr['type'],arr['model_id'],drawing_image,1,arr["device_id"],isHigh,arr["dot"],style_id,arr["is_style_orginal"],arr["platform"],machine))

            sql = "update ap_model_list set use_num =  use_num + 1 where id = %s"
            cur.execute(sql, (arr['model_id']))

            if style_id > 0 :
                sql = "update ap_model_style set use_num =  use_num + 1 where id = %s"
                cur.execute(sql, (style_id))

            conn.commit()

            #if arr['type'] == 3 :
            #    upload_obj = urlparse(drawing_image)
            #    upload_path = upload_obj.path[1:]
            #    ossBucket.delete_object(upload_path)
            #    writeLog("删除用户上传图片：{}".format(upload_path),machine)
                

            if not delay_task :
                r.hset(taskKey,'status',2)
                r.hset(taskKey,'url',url)
                r.hset(taskKey,'prompt',arr['prompt'])
            r.decrby(userLenKey, 1)
            r.expire(userLenKey, 600)
            writeLog("task_id : {} complate!".format(task_id),machine)
        except Exception as e :
            if not delay_task :
                r.hset(taskKey,'status',0)
            r.rpush(queueKey,sdinput)
            writeLog(traceback.format_exc(),machine)
        finally:
            if os.path.exists(pngFile):
                os.remove(pngFile)
            if os.path.exists(webpFile):
                os.remove(webpFile)

    writeLog("loop stop! have a rest!!!",machine)            